package org.example.compressor;

import org.example.tool.HuffmanTree;
import org.example.tool.HuffmanTreeOutputStream;

import java.io.*;
import java.util.HashMap;
import java.util.Map;

public class CompressOutputStream {

    private HuffmanTree huffmanTree;
    private OutputStream outputStream;

    public CompressOutputStream(OutputStream out) {
        outputStream = out;
    }

    private void genMap(Map<Character, String> map, HuffmanTree.Node node, StringBuilder now){
        if(node == null) {
            return;
        }
        if(node.isLeaf()) {
            map.put(node.getChar(), now.toString());
        }
        else {
            genMap(map, node.getZero(), new StringBuilder(now).append('0'));
            genMap(map, node.getOne(), new StringBuilder(now).append('1'));
        }
    }

    public void write(String str) throws IOException {

        new DataOutputStream(outputStream).writeInt(str.length());

        HuffmanTree.FromTextBuilder builder = new HuffmanTree.FromTextBuilder();
        builder.addString(str);
        HuffmanTree tree = builder.build();
        HuffmanTreeOutputStream huffmanTreeWriter = new HuffmanTreeOutputStream(outputStream);
        huffmanTreeWriter.writeTree(tree);

        Map<Character, String> map = new HashMap<>();
        genMap(map, tree.getRoot(), new StringBuilder());

        byte b = 0;
        int bl = 0;
        for(int i = 0; i<str.length(); i++){
            String l = map.get(str.charAt(i));
            for(int j = 0; j<l.length(); j++) {
                b <<= 1;
                if(l.charAt(j) == '1') {
                    b += 1;
                }
                bl ++;
                if (bl == 8) {
                    outputStream.write(b);
                    bl = 0;
                }
            }
        }
        if(bl > 0) {
            for(int i = bl; i < 8; i++) {
                b <<= 1;
            }
            outputStream.write(b);
        }
    }

    public void close() throws IOException {
        this.outputStream.close();
    }

    public void flush() throws IOException {
        this.outputStream.flush();
    }

}
